# The MIT License
#
# Copyright (c) 2011, Ted L. Belser II
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import logging
import traceback
import unittest
import time
import sys
from threading import Thread, Condition, Timer, RLock

class ThreadPoolError(Exception):
    def __init__(self, description, e=None, tb=None):
        self.exceptions = list()
        self.description = description
        if e != None:
            self.AddException(e, tb)

    def __str__(self):
        return "%s: %s" % (self.description, self.exceptions)

    def AddException(self, e, tb):
        self.exceptions.append((e, tb))


class Job:
    def __init__(self, callable=None, args=None, kwargs=None, join_timeout=None):
        self._callable = callable
        self._args = args
        self._kwargs = kwargs
        self._join_timeout = join_timeout
        self._timer = Timer(self._join_timeout, self._TimerExpired)
        self.timed_out = False

    def _TimerExpired(self):
        self.timed_out = True

    def _RunMain(self):
        self._timer.start()
        result = self.Main()
        self._timer.cancel()
        self.ReturnCallback(result)
        return result

    def Main(self):
        return self._callable(*self._args, **self._kwargs)

    def ReturnCallback(self, return_value):
        pass


class ThreadPool:

    class Worker(Thread):
        def __init__(self, pool, logger, logger_level):
            Thread.__init__(self)
            self._pool = pool
            self._logger = logger
            self._logger_level = logger_level
            self._current_job = None

        def __str__(self):
            return self.name

        def run(self):
            self._pool._lock.acquire()
            while True:
                if self._pool._ShuttingDown():
                    self._logger.log(self._logger_level,
                        "Worker %s detected a shutdown signal" % (self.name,))
                    break
                elif self._pool._JobAvailable() == True:
                    self._logger.log(self._logger_level,
                        "A job is available. Worker %s is taking it." % (self.name,))
                    job = self._pool._GetJob()
                    self._pool._JobAccepted(job, self)
                    self._pool._lock.release()
                    self._current_job = job
                    # Execute the job
                    try:
                        result = job._RunMain()
                        job.ReturnCallback(result)
                        self._pool._JobCompleted(job)
                    except Exception, e:
                        tb = traceback.format_exc()
                        self._pool._JobCompleted(job, e, tb)
                    finally:
                        self._current_job = None
                        self._pool._lock.acquire()
                else:
                    self._pool._lock.wait()
                self._logger.log(self._logger_level,
                    "Worker %s notified" % (self.name,))
            self._pool._lock.release()
            self._logger.log(self._logger_level,
                "Worker %s is dying." % (self.name,))

    def __init__(self, max_workers=3, logger_name="pythreadpool",
                 logger_level=logging.INFO):
        self._logger = logging.getLogger(logger_name)
        self._logger.setLevel(logger_level)
        self._logger_level = logger_level
        self._max_workers = max_workers
        # The job_queue list is protected by the condition object
        self._job_queue = list()
        self._jobs_in_progress = dict()
        self._jobs_complete = list()
        self._jobs_failed = dict()
        # Pool state locking mechanisms
        self._lock = Condition(RLock())
        # Shutdown signal
        self._shutting_down = False
        # Create max workers
        self._workers = [ThreadPool.Worker(self, self._logger, self._logger_level) for count in range(self._max_workers)]
        # Start all workers
        [worker.start() for worker in self._workers]

    def _ShuttingDown(self):
        self._lock.acquire()
        result = self._shutting_down
        self._lock.release()
        return result

    def _JobAvailable(self):
        self._lock.acquire()
        if len(self._job_queue) > 0:
            available = True
        else:
            available = False
        self._lock.release()
        return available

    def _GetJob(self):
        self._lock.acquire()
        job = self._job_queue.pop()
        self._lock.release()
        return job

    def _JobAccepted(self, job, worker):
        self._lock.acquire()
        self._jobs_in_progress[job] = worker
        self._lock.release()

    def _JobCompleted(self, job, exception=None, traceback=None):
        self._lock.acquire()
        worker = self._jobs_in_progress.pop(job)
        if exception != None:
            self._jobs_failed[job] = (exception, traceback, worker)
        self._jobs_complete.append(job) 
        self._lock.notifyAll()
        self._lock.release()

    def QueueJob(self, job):
        self._lock.acquire()
        self._job_queue.insert(0, job)
        # Tell waiting workers that the pool's state changed
        # The next available worker will pull from the job queue
        self._lock.notify()
        self._lock.release()        

    def JoinJob(self, job):
        self._lock.acquire()
        if job in self._jobs_complete:
            self._lock.release()
            return
        while job in self._jobs_in_progress:
            self._lock.wait()
        if job in self._jobs_failed:
            (exception, traceback, worker) = self._jobs_failed[job]
            self._lock.release()
            raise ThreadPoolError("Exception in job.", exception, traceback)
        elif job in self._jobs_complete:
            self._lock.release()
            return
        else:
            self._lock.release()
            raise ThreadPoolError("Thread pool has not executed this job.")

    def JoinAllJobs(self, timeout=None):
        self._lock.acquire()
        while len(self._jobs_in_progress) > 0 or len(self._job_queue) > 0:
            self._lock.wait()
        if len(self._jobs_failed) > 0:
            error = ThreadPoolError("At least one job threw an exception.")
            for job, (e, tb, worker) in self._jobs_failed.items():
                error.AddException(e, tb)
            self._lock.release()
            raise error
        else:
            self._lock.release()

    def Shutdown(self, timeout=None):
        self._lock.acquire()
        self._logger.log(self._logger_level, "Sending shutdown signal.")
        self._shutting_down = True
        self._lock.notifyAll()
        self._lock.release()
        #self.JoinAllJobs(timeout)
        for worker in self._workers:
            worker.join()

class ThreadPoolTest(unittest.TestCase):

    class ThreadPoolTestError(Exception):
        pass

    class BasicJob(Job):
        def __init__(self, duration, join_timeout, name=None):
            Job.__init__(self, join_timeout=join_timeout)
            self._duration = duration
            self.name = name
            self.work_complete = False
            self.callback_value = None

        def Main(self):
            timeout = self._duration
            while timeout > 0:
                time.sleep(1)
                print "%s is executing %s." % (self.name, timeout)
                timeout = timeout - 1
            self.work_complete = True
            return True

        def ReturnCallback(self, result):
            self.callback_value = result

    class JobWithException(BasicJob):

        def __init__(self, duration=3, join_timeout=10, name="JobWithException"):
            ThreadPoolTest.BasicJob.__init__(self, duration, join_timeout, name)

        def Main(self):
            timeout = self._duration
            while timeout > 0:
                time.sleep(1)
                print "%s is executing %s." % (self.name, timeout)
                timeout = timeout - 1
            raise ThreadPoolTest.ThreadPoolTestError()


    def SetUp(self):
        pass

    def TearDown(self):
        pass

    def testFewerJobsThanWorkers(self):
        pool = ThreadPool(5)
        job1 = ThreadPoolTest.BasicJob(1, 10, "Job1")
        job2 = ThreadPoolTest.BasicJob(2, 10, "Job2")
        job3 = ThreadPoolTest.BasicJob(3, 10, "Job3")
        pool.QueueJob(job1)
        pool.QueueJob(job2)
        pool.QueueJob(job3)
        pool.JoinAllJobs()
        self.failUnless(job1.work_complete, "Job 1 did not complete")
        self.failUnless(job2.work_complete, "Job 2 did not complete")
        self.failUnless(job3.work_complete, "Job 3 did not complete")
        pool.Shutdown()

    def testFewerWorkersThanJobs(self):
        pool = ThreadPool(3)
        job1 = ThreadPoolTest.BasicJob(1, 10, "Job1")
        job2 = ThreadPoolTest.BasicJob(2, 10, "Job2")
        job3 = ThreadPoolTest.BasicJob(3, 10, "Job3")
        job4 = ThreadPoolTest.BasicJob(2, 10, "Job4")
        job5 = ThreadPoolTest.BasicJob(1, 10, "Job5")
        pool.QueueJob(job1)
        pool.QueueJob(job2)
        pool.QueueJob(job3)
        pool.QueueJob(job4)
        pool.QueueJob(job5)
        pool.JoinAllJobs()
        self.failUnless(job1.work_complete, "Job 1 did not complete")
        self.failUnless(job2.work_complete, "Job 2 did not complete")
        self.failUnless(job3.work_complete, "Job 3 did not complete")
        self.failUnless(job4.work_complete, "Job 4 did not complete")
        self.failUnless(job5.work_complete, "Job 5 did not complete")
        pool.Shutdown()

    def testSingleJobException(self):
        pool = ThreadPool(3)
        job = ThreadPoolTest.JobWithException()
        pool.QueueJob(job)
        try:
            pool.JoinAllJobs()
            self.fail("Job did not raise an exception as expected.")
        except ThreadPoolError as err:
            number_exceptions = len(err.exceptions)
            self.assertTrue(number_exceptions == 1, "Expected one exception, encountered %s" % (number_exceptions))
            self.assertTrue(isinstance(err.exceptions[0][0], ThreadPoolTest.ThreadPoolTestError))
        finally:
            pool.Shutdown()

    def testMultipleJobExceptions(self):
        pool = ThreadPool(4)
        job1 = ThreadPoolTest.JobWithException(name="job1")
        job2 = ThreadPoolTest.JobWithException(name="job2")
        job3 = ThreadPoolTest.JobWithException(name="job3")
        job4 = ThreadPoolTest.BasicJob(3, 10, "job4")

        pool.QueueJob(job1)
        pool.QueueJob(job2)
        pool.QueueJob(job3)
        pool.QueueJob(job4)

        try:
            pool.JoinAllJobs()
            self.fail("Job did not raise an exception as expected.")
        except ThreadPoolError as err:
            self.assertTrue(job4.work_complete, "Job4 did not complete.")
            number_exceptions = len(err.exceptions)
            self.assertTrue(number_exceptions == 3, "Expected one exception, encountered %s" % (number_exceptions))
            self.assertTrue(isinstance(err.exceptions[0][0], ThreadPoolTest.ThreadPoolTestError))
            self.assertTrue(isinstance(err.exceptions[1][0], ThreadPoolTest.ThreadPoolTestError))
            self.assertTrue(isinstance(err.exceptions[2][0], ThreadPoolTest.ThreadPoolTestError))
        finally:
            pool.Shutdown()


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    unittest.main()
